
import argparse, json
from collections import Counter


def generate_json_data(split_path, data_path, max_captions_per_image, min_word_count, max_caption_length):
    split = json.load(open(split_path, 'r'))
    word_count = Counter()

    train_img_paths = []
    train_caption_tokens = []
    validation_img_paths = []
    validation_caption_tokens = []
    test_img_paths = []
    test_caption_tokens = []

    max_length = 0
    for img in split['images']:
        caption_count = 0
        for sentence in img['sentences']:
            if caption_count < max_captions_per_image:
                caption_count += 1
            else:
                break

            try: # support flickr8k datasets.json that doesn't have subfolders
                img['filepath']
            except KeyError:
                filepath_defined = False
            img_path = f"{data_path}/imgs{'/' + img['filepath'] if filepath_defined else ''}/{img['filename']}"

            if img['split'] == 'train':
                train_img_paths.append(img_path)
                train_caption_tokens.append(sentence['tokens'])
            elif img['split'] == 'val':
                validation_img_paths.append(img_path)
                validation_caption_tokens.append(sentence['tokens'])
            elif img['split'] == 'test':
                test_img_paths.append(img_path)
                test_caption_tokens.append(sentence['tokens'])
            max_length = max(max_length, len(sentence['tokens']))
            word_count.update(sentence['tokens'])

    words = [word for word in word_count.keys() if word_count[word] >= min_word_count]
    word_dict = {word: idx + 4 for idx, word in enumerate(words)}
    word_dict['<start>'] = 0
    word_dict['<eos>'] = 1
    word_dict['<unk>'] = 2
    word_dict['<pad>'] = 3

    with open(data_path + '/word_dict.json', 'w') as f:
        json.dump(word_dict, f)

    max_length = min(max_length, max_caption_length)
    train_captions = process_caption_tokens(train_caption_tokens, word_dict, max_length)
    validation_captions = process_caption_tokens(validation_caption_tokens, word_dict, max_length)
    test_captions = process_caption_tokens(test_caption_tokens, word_dict, max_length)

    with open(data_path + '/train_img_paths.json', 'w') as f:
        json.dump(train_img_paths, f)
    with open(data_path + '/val_img_paths.json', 'w') as f:
        json.dump(validation_img_paths, f)
    with open(data_path + '/train_captions.json', 'w') as f:
        json.dump(train_captions, f)
    with open(data_path + '/val_captions.json', 'w') as f:
        json.dump(validation_captions, f)
    with open(data_path + '/test_img_paths.json', 'w') as f:
        json.dump(test_img_paths, f)
    with open(data_path + '/test_captions.json', 'w') as f:
        json.dump(test_captions, f)

def process_caption_tokens(caption_tokens, word_dict, max_length):
    captions = []
    for tokens in caption_tokens:
        tokens = tokens[:max_length]
        token_idxs = [word_dict[token] if token in word_dict else word_dict['<unk>'] for token in tokens]
        captions.append([word_dict['<start>']] + token_idxs + [word_dict['<eos>']] + [word_dict['<pad>']] * (max_length - len(tokens)))

    return captions


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Generate json files')
    parser.add_argument('--split-path', type=str, default='data/coco/dataset.json')
    parser.add_argument('--data-path', type=str, default='data/coco')
    parser.add_argument('--max-captions', type=int, default=5,
                        help='maximum number of captions per image')
    parser.add_argument('--min-word-count', type=int, default=5,
                        help='minimum number of occurences of a word to be included in word dictionary')
    parser.add_argument('--max-caption-length', type=int, default=25,
                        help='maximum number of tokens in a caption')
    args = parser.parse_args()

    generate_json_data(args.split_path, args.data_path, args.max_captions, args.min_word_count, args.max_caption_length)


import json
import torch
from torch.utils.data import Dataset
from collections import defaultdict
from PIL import Image
import json


def pil_loader(path):
    with open(path, 'rb') as f:
        img = Image.open(f)
        return img.convert('RGB')


class ImageCaptionDataset(Dataset):
    def __init__(self, transform, data_path, split_type='train', fraction=1.0, bert=False):
        super(ImageCaptionDataset, self).__init__()
        self.transform = transform
        
        # Load image paths and captions
        img_paths = json.load(open(data_path + f'/{split_type}_img_paths.json', 'r'))
        if bert:
            captions = json.load(open(data_path + f'/{split_type}_captions_bert.json', 'r'))
        else:
            captions = json.load(open(data_path + f'/{split_type}_captions.json', 'r'))

        # Reduce dataset size if fraction is not 1.0
        if fraction != 1.0:
            img_paths = img_paths[:int(len(img_paths) * fraction)]
            captions = captions[:int(len(captions) * fraction)]

        # Preprocess and store data
        self.data = []
        all_captions = defaultdict(list)  # Store all captions for each image path

        for img_path, caption in zip(img_paths, captions):
            img = pil_loader(img_path)
            if self.transform is not None:
                img = self.transform(img)
            self.data.append((torch.FloatTensor(img), torch.tensor(caption)))
            all_captions[img_path].append(caption)

        # Convert all_captions dictionary to a list matching the order of images
        self.all_captions = [all_captions[path] for path in img_paths]

    def __getitem__(self, index):
        img_tensor, caption_tensor = self.data[index]
        all_captions_tensor = torch.tensor(self.all_captions[index])
        return img_tensor, caption_tensor, all_captions_tensor

    def __len__(self):
        return len(self.data)


import json
import spacy
import nltk
from nltk.corpus import stopwords
import pandas as pd
from wordcloud import WordCloud
import matplotlib.pyplot as plt

nltk.download('stopwords')

!python -m spacy download en_core_web_sm
en_core_web_sm = spacy.load("en_core_web_sm")


DATA_PATH = 'data/flickr8k/'
SPLIT_TYPES = ['train', 'val', 'test']

captions = []
for split_type in SPLIT_TYPES:
    captions.extend(json.load(open(DATA_PATH + f'/{split_type}_captions.json', 'r')))

print(f'Loaded {len(captions)} captions')

word_dict = json.load(open(DATA_PATH + '/word_dict.json', 'r'))
vocabulary_size = len(word_dict)

print(f'Vocabulary size: {vocabulary_size}')

# decode captions
decoded_captions = []
for caption in captions:
    decoded_caption = []
    for idx in caption:
        if idx == 0:
            continue
        elif idx == 1:
            break
        else:
            decoded_caption.append(list(word_dict.keys())[list(word_dict.values()).index(idx)])
    decoded_captions.append(decoded_caption)

Loaded 40000 captions
Vocabulary size: 2945


def generate_count_wordcloud(df:pd.DataFrame, top_n_words=30):
    en_stopwords = set(stopwords.words('english'))

    # create custom tokenizer that removes stopwords
    def spacy_tokenizer(text):
        tokens = en_core_web_sm(text)
        return [token for token in tokens if token.text not in en_stopwords]

    tokens = df.text.apply(spacy_tokenizer)
    lowercase_tokens = [token.lower_ for doc in tokens for token in doc]

    # create wordcloud
    wordcloud = WordCloud(
        width=800, height=400, background_color="white", max_words=top_n_words
    ).generate(" ".join(lowercase_tokens))

    # show wordcloud
    plt.figure(figsize=(12, 10))
    plt.imshow(wordcloud, interpolation="bilinear")
    plt.axis("off")
    plt.title(f'Word Cloud Flickr8k Image Captions')
    plt.show()

# create dataframe with captions
decoded_captions_joined = [' '.join(caption) for caption in decoded_captions]
df = pd.DataFrame({'text': decoded_captions_joined})

print(f'Sample decoded caption: {decoded_captions_joined[0]}')

# generate wordcloud
generate_count_wordcloud(df)

Sample decoded caption: a black dog is running after a white dog in the snow


# Distribibution of caption lengths
caption_lengths = [len(caption) for caption in decoded_captions]
df_cl = pd.DataFrame({'caption_length': caption_lengths})

plt.figure()
plt.bar(df_cl.caption_length.value_counts().index, df_cl.caption_length.value_counts())
plt.title('Caption Length Distribution')
plt.xlabel('Caption Length')
plt.ylabel('Count')
plt.show()

VGG(
  (features): Sequential(
    (31): ReLU(inplace=True)
    (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (33): ReLU(inplace=True)
    (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (35): ReLU(inplace=True) <<<<< this is the last layer we use
    (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    ...
  )
)

import torch
import torch.nn as nn
from attention import Attention


class Decoder(nn.Module):
    def __init__(self, vocabulary_size, encoder_dim, tf=False, ado=False, bert=False, attention=False):
        super(Decoder, self).__init__()
        self.use_tf = tf
        self.use_advanced_deep_output = ado
        self.use_bert = bert
        self.use_attention = attention

        # Initializing parameters
        self.encoder_dim = encoder_dim

        # Embeddings
        if bert == True:
            from transformers import BertModel, BertTokenizer
            self.bert_model = BertModel.from_pretrained('bert-base-uncased')
            self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
            self.vocabulary_size = self.bert_model.config.vocab_size
            self.embedding_size = self.bert_model.config.hidden_size # 768

            # Embedding layer using BERT's embeddings
            self.embedding = self.bert_model.get_input_embeddings()

            # Freeze the BERT embeddings
            for param in self.embedding.parameters():
                param.requires_grad = False

            # Delete the BERT model to save memory (and checkpoint size)
            del self.bert_model
        else:
            self.vocabulary_size = vocabulary_size
            self.embedding_size = 512
            self.embedding = nn.Embedding(self.vocabulary_size, self.embedding_size)  # Embedding layer for input words

        # Initial LSTM cell state generators
        self.init_h = nn.Linear(encoder_dim, self.embedding_size)  # For hidden state
        self.init_c = nn.Linear(encoder_dim, self.embedding_size)  # For cell state
        self.tanh = nn.Tanh()

        # Attention mechanism related layers
        self.f_beta = nn.Linear(self.embedding_size, encoder_dim)  # Gating scalar in attention mechanism
        self.sigmoid = nn.Sigmoid()

        # Attention and LSTM components
        self.attention = Attention(encoder_dim, self.embedding_size)  # Attention network
        self.lstm = nn.LSTMCell(self.embedding_size + encoder_dim, self.embedding_size)  # LSTM cell

        # Deep output layers
        if self.use_advanced_deep_output:
            # Advanced DO: Layers for transforming LSTM state, context vector and embedding for DO-RNN
            hidden_dim, intermediate_dim = self.embedding_size, self.embedding_size
            self.f_h = nn.Linear(hidden_dim, intermediate_dim)  # Transforms LSTM hidden state
            self.f_z = nn.Linear(encoder_dim, intermediate_dim)  # Transforms context vector
            self.f_out = nn.Linear(intermediate_dim, self.vocabulary_size)  # Transforms the combined vector (sum of embedding, LSTM state, and context vector) to voc_size
            self.relu = nn.ReLU()  # Activation function
            self.dropout = nn.Dropout()

        # Simple DO: Layer for transforming LSTM state to vocabulary
        self.deep_output = nn.Linear(self.embedding_size, self.vocabulary_size)  # Maps LSTM outputs to vocabulary
        self.dropout = nn.Dropout()

    def forward(self, img_features, captions):
        # Forward pass of the decoder
        batch_size = img_features.size(0)

        # Initialize LSTM state
        h, c = self.get_init_lstm_state(img_features)

        # Teacher forcing setup
        max_timespan = max([len(caption) for caption in captions]) - 1

        if self.use_bert:
            start_token = torch.full((batch_size, 1), self.tokenizer.cls_token_id).long().to(mps_device)
        else:
            start_token = torch.zeros(batch_size, 1).long().to(mps_device)

        # Convert caption tokens to their embeddings
        if self.use_tf:
            caption_embedding = self.embedding(captions)
        else:
            previous_predicted_token_embedding = self.embedding(start_token)

        # Preparing to store predictions and attention weights
        preds = torch.zeros(batch_size, max_timespan, self.vocabulary_size).to(mps_device) # [BATCH_SIZE, TIME_STEPS, VOC_SIZE]
        alphas = torch.zeros(batch_size, max_timespan, img_features.size(1)).to(mps_device) # [BATCH_SIZE, TIME_STEPS, NUM_SPATIAL_FEATURES]

        # Generating captions
        for t in range(max_timespan):
            if self.use_attention:
                context, alpha = self.attention(img_features, h)  # Compute context vector via attention
                gate = self.sigmoid(self.f_beta(h))  # Gating scalar for context
                gated_context = gate * context  # Apply gate to context
            else:
                # If not using attention, treat all parts of the image equally
                alpha = torch.full((batch_size, img_features.size(1)), 1.0 / img_features.size(1), device=mps_device)  # Uniform attention
                context = img_features.mean(dim=1)  # Simply take the mean of the image features
                gated_context = context  # No gating applied

            # Prepare LSTM input
            if self.use_tf:
                lstm_input = torch.cat((caption_embedding[:, t], gated_context), dim=1)  # current embedding + context vector as input vector
            else:
                previous_predicted_token_embedding = previous_predicted_token_embedding.squeeze(1) if previous_predicted_token_embedding.dim() == 3 else previous_predicted_token_embedding
                lstm_input = torch.cat((previous_predicted_token_embedding, gated_context), dim=1)

            # LSTM forward pass
            h, c = self.lstm(lstm_input, (h, c))

            # Generate word prediction
            if self.use_advanced_deep_output:
                # NOTE: could explore alternative positions for dropout
                if self.use_tf:
                    output = self.advanced_deep_output(self.dropout(h), context, caption_embedding[:, t])
                else:
                    output = self.advanced_deep_output(self.dropout(h), context, previous_predicted_token_embedding)
            else:
                output = self.deep_output(self.dropout(h))

            preds[:, t] = output  # Store predictions
            alphas[:, t] = alpha  # Store attention weights

            # Prepare next input word
            if not self.use_tf:
                predicted_token_idxs = output.max(1)[1].reshape(batch_size, 1) # output.max(1)[1] = extract index: [1] of the token with the highest probability: max(1)
                previous_predicted_token_embedding = self.embedding(predicted_token_idxs)

        return preds, alphas

    def get_init_lstm_state(self, img_features):
        # Initializing LSTM state based on image features
        avg_features = img_features.mean(dim=1)

        c = self.init_c(avg_features)  # Cell state
        c = self.tanh(c)

        h = self.init_h(avg_features)  # Hidden state
        h = self.tanh(h)

        return h, c

    def advanced_deep_output(self, h, context, current_embedding):
        # Combine the LSTM state and context vector
        h_transformed = self.relu(self.f_h(h))
        z_transformed = self.relu(self.f_z(context))

        # Sum the transformed vectors with the embedding
        combined = h_transformed + z_transformed + current_embedding

        # Transform the combined vector & compute the output word probability
        return self.relu(self.f_out(combined))

import torch
import torch.nn as nn


class Attention(nn.Module):
    def __init__(self, encoder_dim):
        super(Attention, self).__init__()
        self.U = nn.Linear(512, 512)
        self.W = nn.Linear(encoder_dim, 512)
        self.v = nn.Linear(512, 1)
        self.tanh = nn.Tanh()
        self.softmax = nn.Softmax(1)

    def forward(self, img_features, hidden_state):
        U_h = self.U(hidden_state).unsqueeze(1)
        W_s = self.W(img_features)
        att = self.tanh(W_s + U_h)
        e = self.v(att).squeeze(2)
        alpha = self.softmax(e)
        context = (img_features * alpha.unsqueeze(2)).sum(1)
        return context, alpha

embedding = self.embedding(prev_words)

class EvalMode(Enum):
    VALIDATION = 'val'
    TEST = 'test'

def main(args):
    set_seed(args.seed)
    wandb.init(project='show-attend-and-tell', entity='yvokeller', config=args)

    # ...

    encoder = Encoder(args.network)
    decoder = Decoder(vocabulary_size, encoder.dim, tf=args.tf, ado=args.ado, bert=args.bert, attention=args.attention)

    optimizer = optim.Adam(decoder.parameters(), lr=args.lr)
    scheduler = optim.lr_scheduler.StepLR(optimizer, args.step_size)
    cross_entropy_loss = nn.CrossEntropyLoss().to(mps_device)

    train_loader = DataLoader(
        ImageCaptionDataset(data_transforms, args.data, fraction=args.fraction, bert=args.bert, split_type='train'),
        batch_size=args.batch_size, shuffle=True, num_workers=0, pin_memory=True)

    val_loader = DataLoader(
        ImageCaptionDataset(data_transforms, args.data, fraction=args.fraction, bert=args.bert, split_type='val'),
        batch_size=args.batch_size, shuffle=True, num_workers=0, pin_memory=True)

    test_loader = DataLoader(
        ImageCaptionDataset(data_transforms, args.data, fraction=args.fraction, bert=args.bert, split_type='test'),
        batch_size=args.batch_size, shuffle=True, num_workers=0, pin_memory=True)

    print(f'Starting training with {args}')
    for epoch in range(1, args.epochs + 1):
        train(epoch, encoder, decoder, optimizer, cross_entropy_loss,
              train_loader, word_dict, args.alpha_c, args.log_interval, bert=args.bert, tokenizer=bert_tokenizer, args=args)
        validate(epoch, encoder, decoder, cross_entropy_loss, val_loader,
                 word_dict, args.alpha_c, args.log_interval, bert=args.bert, tokenizer=bert_tokenizer)
        scheduler.step()

    if args.perform_test == True:
        test(epoch, encoder, decoder, cross_entropy_loss, test_loader,
             word_dict, args.alpha_c, args.log_interval, bert=args.bert, tokenizer=bert_tokenizer)

    wandb.finish()


def train(epoch, encoder, decoder, optimizer, cross_entropy_loss, data_loader, word_dict, alpha_c, log_interval, bert=False, tokenizer=None, args={}):
    print(f"Epoch {epoch} - Starting train")

    encoder.eval()
    decoder.train()

    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()
    for batch_idx, (imgs, captions, _) in enumerate(data_loader):
        imgs, captions = Variable(imgs).to(mps_device), Variable(captions).to(mps_device)

        img_features = encoder(imgs)
        optimizer.zero_grad()
        preds, alphas = decoder(img_features, captions)
        targets = captions[:, 1:] # skip <start> token for loss calculation

        # Calculate accuracy
        padding_idx = word_dict['<pad>'] if bert == False else tokenizer.pad_token_id
        acc1 = sequence_accuracy(preds, targets, 1, ignore_index=padding_idx, tokenizer=tokenizer)
        acc5 = sequence_accuracy(preds, targets, 5, ignore_index=padding_idx, tokenizer=tokenizer)

        # Calculate loss
        packed_targets = pack_padded_sequence(targets, [len(tar) - 1 for tar in targets], batch_first=True)[0]
        packed_preds = pack_padded_sequence(preds, [len(pred) - 1 for pred in preds], batch_first=True)[0]

        # encourage total attention (alphas) to be close to 1, thus penalize when sum is far from 1
        att_regularization = alpha_c * ((1 - alphas.sum(1)) ** 2).mean()

        loss = cross_entropy_loss(packed_preds, packed_targets)
        loss += att_regularization # pytorch autograd will calculate gradients for both loss and att_regularization
        loss.backward()
        optimizer.step()

        if bert == True:
            total_caption_length = calculate_caption_lengths(...)
        else:
            total_caption_length = calculate_caption_lengths(...)

        losses.update(loss.item(), total_caption_length)
        top1.update(acc1, total_caption_length)
        top5.update(acc5, total_caption_length)

        if batch_idx % log_interval == 0:
            print(f'Train Batch: [{batch_idx}/{len(data_loader)}]\t'
                  f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
                  f'Top 1 Accuracy {top1.val:.3f} ({top1.avg:.3f})\t'
                  f'Top 5 Accuracy {top5.val:.3f} ({top5.avg:.3f})')

        wandb.log({
            'train_loss': losses.avg, 'train_top1_acc': top1.avg, 'train_top5_acc': top5.avg, 'epoch': epoch,
            'train_loss_raw': losses.val, 'train_top1_acc_raw': top1.val, 'train_top5_acc_raw': top5.val
        })

def validate(epoch, *args, **kwargs):
    print(f"Epoch {epoch} - Starting validation")
    return run_evaluation(epoch, *args, mode=EvalMode.VALIDATION, **kwargs)

def test(epoch, *args, **kwargs):
    print(f"Epoch {epoch} - Starting test")
    return run_evaluation(epoch, *args, mode=EvalMode.TEST, **kwargs)

def run_evaluation(epoch, encoder, decoder, cross_entropy_loss, data_loader, word_dict, alpha_c, log_interval, bert=False, tokenizer=None, mode=EvalMode.VALIDATION):
    encoder.eval()
    decoder.eval()

    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    decoded_captions = [] # list of single assigned caption for each image
    decoded_all_captions = [] # list of list of all captions present in dataset for each image, thus captions may repeat in different lists
    decoded_hypotheses = [] # list of single predicted caption for each image

    predictions_table = wandb.Table(columns=["epoch", "mode", "target_caption", "pred_caption"])

    with torch.no_grad():
        logged_attention_visualizations_count = 0
        for batch_idx, (imgs, captions, all_captions) in enumerate(data_loader):
            imgs, captions = Variable(imgs).to(mps_device), Variable(captions).to(mps_device)
            img_features = encoder(imgs)
            preds, alphas = decoder(img_features, captions)
            targets = captions[:, 1:]

            # Calculate accuracy
            padding_idx = word_dict['<pad>'] if bert == False else tokenizer.pad_token_id
            acc1 = sequence_accuracy(preds, targets, 1, ignore_index=padding_idx, tokenizer=tokenizer)
            acc5 = sequence_accuracy(preds, targets, 5, ignore_index=padding_idx, tokenizer=tokenizer)

            # Calculate loss
            packed_targets = pack_padded_sequence(targets, [len(tar) - 1 for tar in targets], batch_first=True)[0]
            packed_preds = pack_padded_sequence(preds, [len(pred) - 1 for pred in preds], batch_first=True)[0]

            att_regularization = alpha_c * ((1 - alphas.sum(1)) ** 2).mean()

            loss = cross_entropy_loss(packed_preds, packed_targets)
            loss += att_regularization

            if bert == True:
                total_caption_length = calculate_caption_lengths(...)
            else:
                total_caption_length = calculate_caption_lengths(...)

            losses.update(loss.item(), total_caption_length)
            top1.update(acc1, total_caption_length)
            top5.update(acc5, total_caption_length)

            if bert == True:
                # ... DECODE CAPTIONS BERT 
            else:
                # ... DECODE CAPTIONS STANDARD 

            if batch_idx % log_interval == 0:
                print(f'{mode} Batch: [{batch_idx}/{len(data_loader)}]\t'
                    f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
                    f'Top 1 Accuracy {top1.val:.3f} ({top1.avg:.3f})\t'
                    f'Top 5 Accuracy {top5.val:.3f} ({top5.avg:.3f})')

            if mode == EvalMode.TEST:
                # Calculate the start index for the current batch
                batch_start_idx = batch_idx * len(imgs)

                # Log the attention visualizations
                for img_idx, img_tensor in enumerate(imgs):
                    # Skip attention visualization if already logged enough
                    if logged_attention_visualizations_count >= 50:
                        break
                    logged_attention_visualizations_count += 1

                    # Calculate the global index for decoded_hypotheses and decoded_captions lists
                    global_caption_idx = batch_start_idx + img_idx

                    if len(decoded_hypotheses[global_caption_idx]) == 0:
                        print(f'No caption for image {global_caption_idx}, skipping attention visualization')
                        break

                    log_attention_visualization_plot(img_tensor, alphas, decoded_hypotheses, decoded_captions, batch_idx, img_idx, global_caption_idx, encoder)

        bleu_1 = corpus_bleu(decoded_all_captions, decoded_hypotheses, weights=(1, 0, 0, 0))
        bleu_2 = corpus_bleu(decoded_all_captions, decoded_hypotheses, weights=(0.5, 0.5, 0, 0))
        bleu_3 = corpus_bleu(decoded_all_captions, decoded_hypotheses, weights=(0.33, 0.33, 0.33, 0))
        bleu_4 = corpus_bleu(decoded_all_captions, decoded_hypotheses)

        wandb.log({
            'epoch': epoch,
            f'{epoch}_{mode.value}_caption_predictions': predictions_table,
            f'{mode.value}_loss': losses.avg, f'{mode.value}_top1_acc': top1.avg, f'{mode.value}_top5_acc': top5.avg,
            f'{mode.value}_loss_raw': losses.val, f'{mode.value}_top1_acc_raw': top1.val, f'{mode.value}_top5_acc_raw': top5.val,
            f'{mode.value}_bleu1': bleu_1, f'{mode.value}_bleu2': bleu_2, f'{mode.value}_bleu3': bleu_3, f'{mode.value}_bleu4': bleu_4,
        })

        print(f'{mode} Epoch: {epoch}\t'
              f'BLEU-1 ({bleu_1})\t'
              f'BLEU-2 ({bleu_2})\t'
              f'BLEU-3 ({bleu_3})\t'
              f'BLEU-4 ({bleu_4})\t')


from generate_caption import generate_caption_visualization, load_model

import matplotlib.pyplot as plt
%matplotlib inline
%config InlineBackend.figure_format='retina'


WANDB_PROJECT = "yvokeller/show-attend-and-tell/"
WANDB_MODEL_FOLDER = "model/"
WANDB_MODEL_NAME_TEMPLATE = "model_vgg19_X.pth"

def load_model_from_checkpoint(run_id, checkpoint):
    wandb_run = WANDB_PROJECT + run_id
    wandb_model = WANDB_MODEL_FOLDER + WANDB_MODEL_NAME_TEMPLATE.replace("X", str(checkpoint))
    
    return load_model(wandb_run=wandb_run, wandb_model=wandb_model)

def caption_images(img_paths, run_id, checkpoint, figsize=(9, 6), beam_size=3):
    encoder, decoder, bert, model_path, model_config_path = load_model_from_checkpoint(run_id, checkpoint)
    for img_path in img_paths:
        generate_caption_visualization(img_path, encoder, decoder, model_path, model_config_path, beam_size=beam_size, figsize=figsize)


# Image Test Sets
own_images = ['data/mine/train.jpeg', 'data/mine/tashi.jpeg', 'data/mine/lake.jpeg']
test_images_flickr8k = [
    'data/flickr8k/imgs/667626_18933d713e.jpg',
    'data/flickr8k/imgs/280706862_14c30d734a.jpg',
    'data/flickr8k/imgs/3072172967_630e9c69d0.jpg',
    'data/flickr8k/imgs/2654514044_a70a6e2c21.jpg',
    'data/flickr8k/imgs/311146855_0b65fdb169.jpg',
    'data/flickr8k/imgs/2218609886_892dcd6915.jpg',
    'data/flickr8k/imgs/2511019188_ca71775f2d.jpg',
    'data/flickr8k/imgs/2435685480_a79d42e564.jpg',
    'data/flickr8k/imgs/3482062809_3b694322c4.jpg'
]


caption_images(test_images_flickr8k, '8nu0sdou', 8, beam_size=3)


caption_images(test_images_flickr8k, '2asua9lu', 8, beam_size=3)


caption_images(test_images_flickr8k, 'w97u6fc9', 8, beam_size=3)


caption_images(test_images_flickr8k, 'aatq1i5p', 8, beam_size=3)


caption_images(test_images_flickr8k, 'tuzy19bt', 8, beam_size=3)


caption_images(own_images, 'tuzy19bt', 8, beam_size=9)

Criteria	Rating: 1 Point	Rating: 2 Points	Rating: 3 Points	Rating: 4 Points	Rating: 5 Points
Adequacy	Caption does not relate to the image content.	Caption vaguely relates to the image.	Caption covers basic elements in the image.	Caption covers most elements in the image.	Caption accurately covers all key elements of the image.
Semantic Correctness	Caption makes no logical sense.	Caption has significant logical flaws.	Caption is somewhat logical but has some inaccuracies.	Caption is logical with minor inaccuracies.	Caption is completely logical and accurate.
Object Detection	Fails to identify any objects.	Identifies less than 50% of objects correctly.	Identifies about 50% of objects correctly.	Identifies most objects correctly.	Identifies all objects correctly.
Color Detection	Fails to identify any colors.	Identifies less than 50% of colors correctly.	Identifies about 50% of colors correctly.	Identifies most colors correctly.	Identifies all colors correctly.
Attention Visualization	No attention visualization.	Attention visualization is not meaningful.	Attention visualization is somewhat meaningful.	Attention visualization is mostly meaningful.	Attention visualization is very meaningful.

Criteria	Rating (1-5)
Adequacy	4
Semantic Correctness	2
Object Detection	3
Color Detection	4
Attention Visualization	4

Criteria	Rating (1-5)
Adequacy	3
Semantic Correctness	2
Object Detection	2
Color Detection	3
Attention Visualization	N/A

Criteria	Rating (1-5)
Adequacy	2
Semantic Correctness	1
Object Detection	2
Color Detection	2
Attention Visualization	2

Criteria	Rating (1-5)
Adequacy	1
Semantic Correctness	1
Object Detection	1
Color Detection	3
Attention Visualization	N/A

Introduction¶

Data Loader¶

Preparing Dataset for the "Show, Attend and Tell" Model¶

Karpathy's Data Splits for Comparison to Original Paper¶

Role of Tokenization and Special Tokens¶

Maximum Caption Length and Reasoning¶

Preparing Data for Faster Loading During Training¶

Implementation Details¶

Code¶

ImageCaptionDataset¶

Utilizing Pinned Memory¶

Preprocessing Done Once Before Training¶

Key Features of the ImageCaptionDataset¶

Code¶

Explorative Data Analysis¶

Captions¶

Images¶

Understanding Show, Attend and Tell¶

The Encoder¶

Code¶

The Decoder¶

Code¶

Attention Mechanisms¶

Stochastic Hard Attention¶

1. Attention Location Representation¶

2. Context Vector Computation¶

3. Attention as Multinoulli Distribution¶

4. Objective Function (Variational Lower Bound)¶

5. Gradient Approximation via Monte Carlo Sampling¶

6. Variance Reduction Techniques¶

7. Final Learning Rule¶

Connection to REINFORCE Learning Rule¶

Deterministic Soft Attention¶

1. Expectation of the Context Vector¶

2. Deterministic Attention Model¶

3. Normalized Weighted Geometric Mean (NWGM)¶

4. Simplification for Learning¶

5. Approximation to Marginal Likelihood¶

Conclusion¶

Implementation of the Attention Mechanism¶

Calculating Attention Weights and Context Vector¶

Changes to the Implementation¶

Deep Output Layer¶

Explanation¶

Implementing the Deep Output Layer¶

Integration of BERT Embeddings in the Decoder¶

BERT Data Preprocessing¶

BERT Embedding Integration¶

Decoder Architecture Adjustments¶

Considerations for BERT Integration¶

Metrics¶

BLEU Score¶

Overview of BLEU Score¶

Formula¶

Application¶

Top-N Accuracy¶

Understanding Top-N Accuracy¶

Formula¶

Application¶

Experiment Setup¶

Overview¶

Tracking and Logging with Weights & Biases (W&B)¶

Dataset and Splits¶

Model Variants¶

Controllable Hyperparameters¶

Experiment Focus¶

Training¶

Train, Validation, and Test Modes¶

Key Components of the Training Procedure¶

Loss Calculation and Regularization¶

Accuracy and BLEU Score Calculation¶

Logging and Visualization¶

Code¶

Results¶

Plot 1: Average Training vs. Validation Loss¶

Plot 2: Accuracy¶

Plot 3: BLEU Scores¶

Overall Takeaways from Chart Analysis¶

Proposed Next Steps¶

Bonus Models¶